import os
import socket
import numpy as np
import cv2
from datetime import datetime

'''
以下是一个简单的 TCP 接收端 Python 方法，用于接收来自客户端（例如工业相机发送的图像数据）的数据。接收端将不断监听连接，并将接收到的数据按图片格式解码并保存。

假设你通过 TCP 接收的是一张张的图像数据流（例如 Bayer 格式的二进制图像数据），以下代码可以帮助你解码并保存这些图片。

主要用来监听GrabImage.py抓拍程序，并传输相机采集得图像数据。
'''

def encoder_to_image(img_data):
    # 将字节数据转换为 NumPy 数组并解码为图像
    np_array = np.frombuffer(img_data, np.uint8)
    img = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
    return img

def receive_image_over_tcp(server_ip, server_port, save_folder):
    # 创建 TCP 服务端 Socket
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind((server_ip, server_port))
    server_socket.listen(5)
    print(f"Listening on {server_ip}:{server_port}...")

    while True:
        client_socket, _ = server_socket.accept()
        print("Client connected.")

        # 接收图像数据
        data = b""
        while True:
            packet = client_socket.recv(4096)  # 每次接收 4096 字节
            if not packet:
                break
            data += packet

        # 关闭客户端连接
        client_socket.close()

        img = encoder_to_image(data)
        # 保存图像，以时间戳命名
        if img is not None:
            image = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            filename = f"{save_folder}/received_image_{timestamp}.jpg"
            cv2.imwrite(filename, image)
            print(f"Image saved as {filename}")
        else:
            print("Failed to decode image.")

            return

if __name__ == '__main__':
    # 使用示例
    save_folder = '/ai_data1/wirerope/hkData/gongzhuang/received_images2/'
    os.makedirs(save_folder,exist_ok=True)
    server_ip = '0.0.0.0'
    server_port = 8084
    receive_image_over_tcp(server_ip, server_port, save_folder)  # 指定存储目录


